import requests
from pymilvus import connections, Collection
import re

# ANSI 转义码定义
HIGHLIGHT = "\033[1;31m"  # 红色高亮
RESET = "\033[0m"         # 重置颜色

# 连接 Milvus
connections.connect(host='localhost', port='19530')
collection = Collection(name="text_collection")
collection.load()

# Xinference 模型 UID
model_uid = "bge-small-en-v1.5-a5JDNlUy"  # 使用你的实际 model_uid
embed_url = "http://localhost:9997/v1/embeddings"

def highlight_match(text, query):
    # 不区分大小写地高亮匹配的查询词
    pattern = re.compile(re.escape(query), re.IGNORECASE)
    highlighted = pattern.sub(f"{HIGHLIGHT}\\g<0>{RESET}", text)
    return highlighted

def search_query(query_text):
    payload = {"model": model_uid, "input": query_text}
    query_embedding = requests.post(embed_url, json=payload).json()["data"][0]["embedding"]
    results = collection.search(
        data=[query_embedding],
        anns_field="embedding",
        param={"metric_type": "COSINE", "params": {"nprobe": 10}},
        limit=5,
        output_fields=["text"]
    )
    threshold = 0.7  # 设置相似度阈值
    found = False
    for result in results[0]:
        similarity = result.distance
        if similarity >= threshold:
            text = result.entity.get("text")
            highlighted_text = highlight_match(text, query_text)
            print(f"Similarity: {similarity:.4f}, Text: {highlighted_text}...")
            found = True
    if not found:
        print(f"没有找到相似度高于 {threshold} 的结果")

# 交互式查询
while True:
    query = input("请输入查询词（输入 'exit' 退出）：")
    if query.lower() == "exit":
        break
    search_query(query)